module aru_reduce_transpose (
    input logic                 clk,
    input logic                 rst_n,
          aru_reduce_ctrl_if.in u_aru_ctrl_if,
          aru_payload_if.in     u_aru_payload_if,
          aru_reduce_pld_if.out u_aru_reduce_pld_if
);
    // ===== 暂存控制信号 =====
    logic reduce_m_reg;

    logic lst_req_in_instr;
    logic ctrl_rdy, ctrl_vld;

    // 控制信号握手时暂存
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            ctrl_rdy     <= 1'b1;
            reduce_m_reg <= 1'b0;
        end else if (ctrl_rdy == 1'b0) begin
            if (u_aru_reduce_pld_if.vld && u_aru_reduce_pld_if.rdy && u_aru_reduce_pld_if.sdb.eom && u_aru_reduce_pld_if.sdb.eon) begin
                ctrl_rdy <= 1'b1;
            end
        end else begin
            if (u_aru_ctrl_if.vld && u_aru_ctrl_if.rdy) begin
                ctrl_rdy     <= 1'b0;
                reduce_m_reg <= u_aru_ctrl_if.reduce_m;
            end
        end
    end

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            ctrl_vld <= 1'b0;
        end else if (ctrl_vld == 1'b0) begin
            if (u_aru_ctrl_if.vld && u_aru_ctrl_if.rdy) begin
                ctrl_vld <= 1'b1;
            end
        end else begin
            if (u_aru_payload_if.vld && u_aru_payload_if.rdy && u_aru_payload_if.sdb.eom && u_aru_payload_if.sdb.eon) begin
                ctrl_vld <= 1'b0;
            end
        end
    end

    assign u_aru_ctrl_if.rdy = ctrl_rdy;

    // ===== 数据处理逻辑使用暂存的控制信号 =====
    logic            [`P_ARU*`N0-1:0] msk;
    logic            [`P_ARU*`N0-1:0] transposed_msk;
    aru_reduce_dat_t                  transposed_dat;

    genvar n0, m0;
    generate
        for (n0 = 0; n0 < `N0; n0 = n0 + 1) begin : gen_n0
            for (m0 = 0; m0 < `P_ARU; m0 = m0 + 1) begin : gen_m0
                assign msk[m0*`N0+n0] = (m0 < u_aru_payload_if.sdb.vld_m) && (n0 < u_aru_payload_if.sdb.vld_n);
                assign transposed_msk[n0*`P_ARU+m0] = msk[m0*`N0+n0];
                assign transposed_dat[n0*`P_ARU+m0].bf16.dat = u_aru_payload_if.dat.dat[m0*`N0+n0];
                assign transposed_dat[n0*`P_ARU+m0].bf16.rsv = 16'd0;
            end
        end
    endgenerate

    // Stage 1 - 使用暂存的控制信号
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            u_aru_reduce_pld_if.vld <= 1'b0;
            u_aru_reduce_pld_if.dat <= 'd0;
            u_aru_reduce_pld_if.sdb <= 'd0;
            u_aru_reduce_pld_if.msk <= 'd0;
        end else begin
            if (reduce_m_reg) begin  // 使用暂存的值
                u_aru_reduce_pld_if.dat <= transposed_dat;
            end else begin
                for (int i = 0; i < `P_ARU * `N0; i++) begin
                    u_aru_reduce_pld_if.dat[i].bf16.dat <= u_aru_payload_if.dat.dat[i];
                    u_aru_reduce_pld_if.dat[i].bf16.rsv <= 16'd0;
                end
            end
            u_aru_reduce_pld_if.vld <= u_aru_payload_if.vld;
            u_aru_reduce_pld_if.sdb <= u_aru_payload_if.sdb;
            u_aru_reduce_pld_if.msk <= reduce_m_reg ? transposed_msk : msk;  // 使用暂存的值
        end
    end

    assign u_aru_payload_if.rdy = u_aru_reduce_pld_if.rdy & ctrl_vld;

endmodule
